
#include <CL/cl.h>
#include <iostream>
#include <fstream>
#include <vector>
#include <cmath>
#include <cassert>

#define CHECK_CL(err, msg) \
    if (err != CL_SUCCESS) { \
        std::cerr << "❌ OpenCL error " << err << " at " << msg << std::endl; \
        exit(1); \
    }



std::string loadKernelSource(const std::string& filename) {
    std::string kernel_path = KERNELS_DIR "/" + filename;
    std::ifstream file(kernel_path);
    if (!file.is_open()) {
        throw std::runtime_error("Failed to open kernel file: " + kernel_path);
    }
    return std::string((std::istreambuf_iterator<char>(file)), 
                      std::istreambuf_iterator<char>());
}


int main() {
    cl_int err;

    // 1️⃣ 选择第一个平台和设备
    cl_uint num_platforms = 0;
    CHECK_CL(clGetPlatformIDs(0, nullptr, &num_platforms), "get platforms");
    std::vector<cl_platform_id> platforms(num_platforms);
    CHECK_CL(clGetPlatformIDs(num_platforms, platforms.data(), nullptr), "get platform list");

    cl_platform_id platform = platforms[0];
    cl_uint num_devices = 0;
    CHECK_CL(clGetDeviceIDs(platform, CL_DEVICE_TYPE_ALL, 0, nullptr, &num_devices), "get device count");

    std::vector<cl_device_id> devices(num_devices);
    CHECK_CL(clGetDeviceIDs(platform, CL_DEVICE_TYPE_ALL, num_devices, devices.data(), nullptr), "get devices");
    cl_device_id device = devices[0];

    // 打印设备名
    char device_name[256];
    clGetDeviceInfo(device, CL_DEVICE_NAME, sizeof(device_name), device_name, nullptr);
    std::cout << "Using device: " << device_name << std::endl;

    // 2️⃣ 创建上下文和命令队列
    cl_context context = clCreateContext(nullptr, 1, &device, nullptr, nullptr, &err);
    CHECK_CL(err, "create context");
    cl_command_queue queue = clCreateCommandQueue(context, device, 0, &err);
    CHECK_CL(err, "create queue");

    // 3️⃣ 读取并编译 kernel
    std::string src = loadKernelSource("flash_attn_f32.cl");
    const char* src_c = src.c_str();
    size_t src_len = src.size();
    cl_program program = clCreateProgramWithSource(context, 1, &src_c, &src_len, &err);
    CHECK_CL(err, "create program");

    // 这里你需要定义 BLOCK_M、BLOCK_N、DK、DV 宏
    const char* build_opts = "-DBLOCK_M=2 -DBLOCK_N=2 -DDK=4 -DDV=4";
    err = clBuildProgram(program, 1, &device, build_opts, nullptr, nullptr);
    if (err != CL_SUCCESS) {
        size_t log_size;
        clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, 0, nullptr, &log_size);
        std::vector<char> log(log_size);
        clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, log_size, log.data(), nullptr);
        std::cerr << "❌ Build error:\n" << log.data() << std::endl;
        return 1;
    }

    // 4️⃣ 创建 kernel
    cl_kernel kernel_flash = clCreateKernel(program, "flash_attn_f32", &err);
    CHECK_CL(err, "create kernel flash_attn_f32");
    cl_kernel kernel_flash_q1 = clCreateKernel(program, "flash_attn_f32_q1", &err);
    CHECK_CL(err, "create kernel flash_attn_f32_q1");

    // 5️⃣ 准备测试数据
    const int n_q = 2, n_kv = 2, n_head = 1, n_head_kv = 1;
    const float scale = 1.0f / std::sqrt(4.0f);

    // Q, K, V 各有 n_q * DK, n_kv * DK, n_kv * DV
    std::vector<float> Q = {0.1f, 0.2f, 0.3f, 0.4f,
                            0.5f, 0.6f, 0.7f, 0.8f};
    std::vector<float> K = {0.2f, 0.3f, 0.4f, 0.5f,
                            0.6f, 0.7f, 0.8f, 0.9f};
    std::vector<float> V = {0.3f, 0.4f, 0.5f, 0.6f,
                            0.7f, 0.8f, 0.9f, 1.0f};
    std::vector<float> O(8, 0.0f);

    // 6️⃣ 创建缓冲区
    cl_mem q_buf = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, sizeof(float)*Q.size(), Q.data(), &err);
    CHECK_CL(err, "q buffer");
    cl_mem k_buf = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, sizeof(float)*K.size(), K.data(), &err);
    CHECK_CL(err, "k buffer");
    cl_mem v_buf = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, sizeof(float)*V.size(), V.data(), &err);
    CHECK_CL(err, "v buffer");
    cl_mem o_buf = clCreateBuffer(context, CL_MEM_READ_WRITE, sizeof(float)*O.size(), nullptr, &err);
    CHECK_CL(err, "o buffer");

    // 7️⃣ 设置 kernel 参数（简化版：不使用 mask、sink）
    size_t arg = 0;
    cl_ulong zero64 = 0;
    int zero = 0;

    // --- 替换用的 set_common_args（严格匹配 kernel 的 40 个参数） ---
    auto set_common_args = [&](cl_kernel kernel) {
        cl_int err;
        size_t arg = 0;
    
        // 准备常量值（按类型精确）
        cl_ulong zero64 = 0;
        cl_mem null_mem = (cl_mem)0;      // 用于 mask_void / sinks_void = NULL
        float max_bias = 0.0f;
        float m0 = 0.0f;
        float m1 = 0.0f;
        int n_head_log2 = 0;
        float logit_softcap = 0.0f;
        int mask_ne2 = 0;
        int mask_ne3 = 0;
        cl_ulong mask_offset = 0;
        cl_ulong mask_nb1 = 0, mask_nb2 = 0, mask_nb3 = 0;
        cl_ulong sinks_offset = 0;
    
        // 0 const global void * q_void
        err = clSetKernelArg(kernel, arg++, sizeof(cl_mem), &q_buf); CHECK_CL(err, "set arg q_void");
        // 1 ulong q_offset
        err = clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero64); CHECK_CL(err, "set arg q_offset");
        // 2 const global void * k_void
        err = clSetKernelArg(kernel, arg++, sizeof(cl_mem), &k_buf); CHECK_CL(err, "set arg k_void");
        // 3 ulong k_offset
        err = clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero64); CHECK_CL(err, "set arg k_offset");
        // 4 const global void * v_void
        err = clSetKernelArg(kernel, arg++, sizeof(cl_mem), &v_buf); CHECK_CL(err, "set arg v_void");
        // 5 ulong v_offset
        err = clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero64); CHECK_CL(err, "set arg v_offset");
        // 6 global void * o_void
        err = clSetKernelArg(kernel, arg++, sizeof(cl_mem), &o_buf); CHECK_CL(err, "set arg o_void");
        // 7 ulong o_offset
        err = clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero64); CHECK_CL(err, "set arg o_offset");
    
        // 8 const float scale
        err = clSetKernelArg(kernel, arg++, sizeof(float), &scale); CHECK_CL(err, "set arg scale");
        // 9 const int n_q
        err = clSetKernelArg(kernel, arg++, sizeof(int), &n_q); CHECK_CL(err, "set arg n_q");
        // 10 const int n_kv
        err = clSetKernelArg(kernel, arg++, sizeof(int), &n_kv); CHECK_CL(err, "set arg n_kv");
        // 11 const int is_causal
        err = clSetKernelArg(kernel, arg++, sizeof(int), &zero); CHECK_CL(err, "set arg is_causal");
        // 12 const int n_head
        err = clSetKernelArg(kernel, arg++, sizeof(int), &n_head); CHECK_CL(err, "set arg n_head");
    
        // 13..15 q_nb1, q_nb2, q_nb3 (ulong)
        err = clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero64); CHECK_CL(err, "set arg q_nb1");
        err = clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero64); CHECK_CL(err, "set arg q_nb2");
        err = clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero64); CHECK_CL(err, "set arg q_nb3");
    
        // 16..18 k_nb1, k_nb2, k_nb3
        err = clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero64); CHECK_CL(err, "set arg k_nb1");
        err = clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero64); CHECK_CL(err, "set arg k_nb2");
        err = clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero64); CHECK_CL(err, "set arg k_nb3");
    
        // 19..21 v_nb1, v_nb2, v_nb3
        err = clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero64); CHECK_CL(err, "set arg v_nb1");
        err = clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero64); CHECK_CL(err, "set arg v_nb2");
        err = clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero64); CHECK_CL(err, "set arg v_nb3");
    
        // 22..24 o_nb1, o_nb2, o_nb3
        err = clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero64); CHECK_CL(err, "set arg o_nb1");
        err = clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero64); CHECK_CL(err, "set arg o_nb2");
        err = clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero64); CHECK_CL(err, "set arg o_nb3");
    
        // 25..27 floats: max_bias, m0, m1
        err = clSetKernelArg(kernel, arg++, sizeof(float), &max_bias); CHECK_CL(err, "set arg max_bias");
        err = clSetKernelArg(kernel, arg++, sizeof(float), &m0); CHECK_CL(err, "set arg m0");
        err = clSetKernelArg(kernel, arg++, sizeof(float), &m1); CHECK_CL(err, "set arg m1");
    
        // 28 int n_head_log2
        err = clSetKernelArg(kernel, arg++, sizeof(int), &n_head_log2); CHECK_CL(err, "set arg n_head_log2");
    
        // 29 float logit_softcap
        err = clSetKernelArg(kernel, arg++, sizeof(float), &logit_softcap); CHECK_CL(err, "set arg logit_softcap");
    
        // 30 int n_head_kv
        err = clSetKernelArg(kernel, arg++, sizeof(int), &n_head_kv); CHECK_CL(err, "set arg n_head_kv");
    
        // 31 const global void* mask_void  (pass NULL)
        err = clSetKernelArg(kernel, arg++, sizeof(cl_mem), &null_mem); CHECK_CL(err, "set arg mask_void");
        // 32 ulong mask_offset
        err = clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &mask_offset); CHECK_CL(err, "set arg mask_offset");
        // 33..35 mask_nb1, mask_nb2, mask_nb3
        err = clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &mask_nb1); CHECK_CL(err, "set arg mask_nb1");
        err = clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &mask_nb2); CHECK_CL(err, "set arg mask_nb2");
        err = clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &mask_nb3); CHECK_CL(err, "set arg mask_nb3");
    
        // 36..37 mask_ne2, mask_ne3  (int)
        err = clSetKernelArg(kernel, arg++, sizeof(int), &mask_ne2); CHECK_CL(err, "set arg mask_ne2");
        err = clSetKernelArg(kernel, arg++, sizeof(int), &mask_ne3); CHECK_CL(err, "set arg mask_ne3");
    
        // 38 const global void* sinks_void (pass NULL)
        err = clSetKernelArg(kernel, arg++, sizeof(cl_mem), &null_mem); CHECK_CL(err, "set arg sinks_void");
        // 39 ulong sinks_offset
        err = clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &sinks_offset); CHECK_CL(err, "set arg sinks_offset");
    
        // 最后验证我们确实设置到了 40 个参数
        std::cout << "set_common_args: set " << arg << " kernel args\n";
    };


    set_common_args(kernel_flash);
    set_common_args(kernel_flash_q1);

    cl_uint num_args;
    clGetKernelInfo(kernel_flash, CL_KERNEL_NUM_ARGS, sizeof(num_args), &num_args, nullptr);
    std::cout << "Kernel expects " << num_args << " args" << std::endl;
    std::cout << "We set up to arg index: " << arg << std::endl;


    // 8️⃣ 运行 kernel
    size_t global_work_size1[2] = {2, 1};
    size_t local_work_size1[2] = {2, 1};

    std::cout << "Running flash_attn_f32..." << std::endl;
    err = clEnqueueNDRangeKernel(queue, kernel_flash, 2, nullptr, global_work_size1, nullptr, 0, nullptr, nullptr);
    CHECK_CL(err, "enqueue flash_attn_f32");
    clFinish(queue);

    err = clEnqueueReadBuffer(queue, o_buf, CL_TRUE, 0, sizeof(float)*O.size(), O.data(), 0, nullptr, nullptr);
    CHECK_CL(err, "read O");

    std::cout << "Output from flash_attn_f32:\n";
    for (float v : O) std::cout << v << " ";
    std::cout << std::endl;

    // 清空输出缓冲区后再测另一个 kernel
    std::fill(O.begin(), O.end(), 0.0f);
    clEnqueueWriteBuffer(queue, o_buf, CL_TRUE, 0, sizeof(float)*O.size(), O.data(), 0, nullptr, nullptr);

    std::cout << "Running flash_attn_f32_q1..." << std::endl;
    err = clEnqueueNDRangeKernel(queue, kernel_flash_q1, 2, nullptr, global_work_size1, nullptr, 0, nullptr, nullptr);
    CHECK_CL(err, "enqueue flash_attn_f32_q1");
    clFinish(queue);

    err = clEnqueueReadBuffer(queue, o_buf, CL_TRUE, 0, sizeof(float)*O.size(), O.data(), 0, nullptr, nullptr);
    CHECK_CL(err, "read O q1");

    std::cout << "Output from flash_attn_f32_q1:\n";
    for (float v : O) std::cout << v << " ";
    std::cout << std::endl;

    // 9️⃣ 清理
    clReleaseKernel(kernel_flash);
    clReleaseKernel(kernel_flash_q1);
    clReleaseProgram(program);
    clReleaseMemObject(q_buf);
    clReleaseMemObject(k_buf);
    clReleaseMemObject(v_buf);
    clReleaseMemObject(o_buf);
    clReleaseCommandQueue(queue);
    clReleaseContext(context);

    std::cout << "✅ Done.\n";
    return 0;
}
